package nsfs

import (
	"GLibHac/fssystem"
	"crypto/aes"
)

type Aes128CtrStorage struct {
	*SectorStorage
	decrypt       *Aes128CtrTransform
	counterOffset int64
	Counter       []byte
}

func (i *Aes128CtrStorage) Read(position int64, destination []byte) (int, error) {
	n, err := i.SectorStorage.Read(position, destination)
	if err != nil {
		return 0, err
	}
	i.updateCounter(i.counterOffset + position)
	i.decrypt.TransformBlock(destination)
	return n, nil
}

func (i *Aes128CtrStorage) Size() int64 {
	return i.SectorStorage.Size()
}

func (i *Aes128CtrStorage) updateCounter(offset int64) {
	off := uint64(offset) >> 4
	for j := uint(0); j < 0x7; j++ {
		i.Counter[0x10-j-1] = (byte)(off & 0xFF)
		off >>= 8
	}

	// Because the value stored in the counter is offset >> 4, the top 4 bits
	// of byte 8 need to have their original value preserved
	i.Counter[8] = (byte)(int(i.Counter[8]&0xF0) | (int)(off&0x0F))
}

func NewAes128CtrStorage(baseStorage fssystem.IStorage, key []byte, counterOffset int64, counterHi []byte) (r *Aes128CtrStorage, err error) {
	bytes := make([]byte, Aes128CtrTransformBlockSizeBytes)
	copy(bytes, counterHi[:8])
	transform, err := NewAes128CtrTransform(key, bytes)
	if err != nil {
		return nil, err
	}
	r = &Aes128CtrStorage{
		SectorStorage: NewSectorStorage(baseStorage, aes.BlockSize),
		counterOffset: counterOffset,
		decrypt:       transform,
		Counter:       transform.Counter,
	}
	r.ReadFunc = r.Read
	r.SizeFunc = r.Size
	return
}
